import random
import cv2
from matplotlib import pyplot as plt
import xml.etree.ElementTree as ET
import albumentations as A
import os
import time
import glob
from tqdm import trange

BOX_COLOR = (255, 0, 0)  # Red
TEXT_COLOR = (255, 255, 255)  # White
# 生成图片的数量倍数
GENERATED_PICS_SIZE = 10


def visualize_bbox(img, bbox, class_name, color=BOX_COLOR, thickness=2):
    """Visualizes a single bounding box on the image"""
    # x_min, y_min, w, h = bbox
    # x_min, x_max, y_min, y_max = int(x_min), int(x_min + w), int(y_min), int(y_min + h)
    x_min, y_min, x_max, y_max = bbox
    print(x_min, y_min, x_max, y_max)

    cv2.rectangle(img, (int(x_min), int(y_min)), (int(x_max), int(y_max)),
                  color=color, thickness=thickness)

    ((text_width, text_height), _) = cv2.getTextSize(class_name,
                                                     cv2.FONT_HERSHEY_SIMPLEX,
                                                     0.35, 1)
    cv2.rectangle(img, (int(x_min), int(y_min) - int(1.3 * text_height)),
                  (int(x_min) + text_width, int(y_min)), BOX_COLOR, -1)
    cv2.putText(
        img,
        text=class_name,
        org=(int(x_min), int(y_min) - int(0.3 * text_height)),
        fontFace=cv2.FONT_HERSHEY_SIMPLEX,
        fontScale=0.35,
        color=TEXT_COLOR,
        lineType=cv2.LINE_AA,
    )
    return img


def visualize(image, bboxes, category_ids, category_id_to_name):
    img = image.copy()
    for bbox, category_id in zip(bboxes, category_ids):
        class_name = category_id_to_name[category_id]
        img = visualize_bbox(img, bbox, class_name)
    plt.axis('off')
    plt.imshow(img)
    plt.show()


def saveNewAnnotation(new_xml_path, new_jpg_path, xml_path, bboxes):
    in_file = open(xml_path, encoding='utf-8')
    new_file = in_file
    tree = ET.parse(new_file)
    root = tree.getroot()
    dir,filename = os.path.split(new_jpg_path)

    root[0].text = os.path.basename(dir)

    root[1].text = filename

    root[2].text = new_jpg_path
    idx = 0
    l1 = len(list(root.iter('object')))
    l2 = len(bboxes)
    if l1 != l2:
        print(f"{l1} != {l2}, {filename}")
    for idx,obj in enumerate(root.iter('object')):
        if idx < len(bboxes):
            obj[4][0].text = str(round(bboxes[idx][0]))
            obj[4][1].text = str(round(bboxes[idx][1]))
            obj[4][2].text = str(round(bboxes[idx][2]))
            obj[4][3].text = str(round(bboxes[idx][3]))
        else:
            root.remove(obj)
        
    tree.write(new_xml_path, 'UTF-8')


def getAnnotation(xml_path):
    '''
    :param xml_path:
    :return: bboxes, category_ids
    '''
    in_file = open(xml_path, encoding='utf-8')
    try:
        tree = ET.parse(in_file)
    except:
        return [], []
    root = tree.getroot()

    bboxes = []
    category_ids = []

    for obj in root.iter('object'):
        cls = obj.find('name').text

        xmlbox = obj.find('bndbox')
        bbox = [int(float(xmlbox.find('xmin').text)),
                int(float(xmlbox.find('ymin').text)),
                int(float(xmlbox.find('xmax').text)),
                int(float(xmlbox.find('ymax').text))]
        bboxes.append(bbox)
        category_ids.append(cls)
    return bboxes, category_ids


def process_aug(pic_dir,xml_dir,out_pic_dir,out_xml_dir,aug_times):
    PICS_DIR = pic_dir
    XML_DIR = xml_dir
    print("picsdir:",PICS_DIR,",XMLDIR:",XML_DIR)
    new_jpg_path_prefix = out_pic_dir
    new_xml_path_prefix = out_xml_dir
    if not os.path.exists(new_jpg_path_prefix):
        print("make dir:",new_jpg_path_prefix)
        os.makedirs(new_jpg_path_prefix)
    if not os.path.exists(new_xml_path_prefix):
        print("make dir:",new_xml_path_prefix)
        os.makedirs(new_xml_path_prefix)
    pic_paths = glob.glob(os.path.join(PICS_DIR, '*.jpg'))

    
    for i in trange(len(pic_paths)):
        jpg_path = pic_paths[i]

        (path, filename) = os.path.split(jpg_path)
        filename_nosuffix = filename.split('.')[0]
        print("filename_nosuffix:",filename_nosuffix)

        xml_path = os.path.join(XML_DIR,filename_nosuffix+".xml")
        print("jpg_path:",jpg_path,",xml_path:",xml_path)
        # print(jpg_path.split('.'))
        for i in range(aug_times):
            image = cv2.imread(jpg_path)
            # print(width, ", ", height)
            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
            new_jpg_path = os.path.join(new_jpg_path_prefix, filename_nosuffix + "_aug_" + str(i) + ".jpg")
            new_xml_path = os.path.join(new_xml_path_prefix,filename_nosuffix + "_aug_" + str(i) + ".xml")
            bboxes, category_ids = getAnnotation(xml_path=xml_path)
            if len(bboxes) == 0 & len(category_ids) == 0:
                continue
            category_id_to_name = {}
            for i in range(len(category_ids)):
                category_id_to_name[category_ids[i]] = category_ids[i]
            # 变换操作
            # 水平反转，高斯模糊，gamma变换，亮度变化，
            transform = A.Compose(
                [
                    A.HorizontalFlip(p=0.5),
                    A.Rotate(limit=70, p=0.3),
                    A.ShiftScaleRotate(shift_limit=0.0625,scale_limit=0, rotate_limit=0,p=0.3),
                    A.GaussianBlur(blur_limit=1, p=0.5),
                    A.ColorJitter(brightness=0.05, contrast=0.05,
                                  saturation=0.02,
                                  hue=0.02, always_apply=False, p=1)
                ],
                bbox_params=A.BboxParams(format='pascal_voc',
                                         label_fields=['category_ids']),
            )
            transformed = transform(image=image, bboxes=bboxes,
                                    category_ids=category_ids)
            image = transformed['image']
            bboxes = transformed['bboxes']
            category_ids = transformed['category_ids']
            # print(bboxes)
            image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
            cv2.imencode('.jpg', image)[1].tofile( new_jpg_path)
            #visualize(image, bboxes, category_ids, category_id_to_name)
            saveNewAnnotation(new_xml_path, new_jpg_path, xml_path, bboxes)
        time.sleep(1)




if __name__ == '__main__':

    pic_dir = "/home/cxz/yolo/data/train/images"#原始图像文件夹
    xml_dir = "/home/cxz/yolo/data/train/labels"#原始xml文件夹
    aug_pic_dir = "/home/cxz/yolo/data/train/new_images"#增强结果图像文件夹
    aug_xml_dir = "/home/cxz/yolo/data/train/new_labels"#增强结果xml文件夹
    aug_times = 15 #每张图片增强多少倍
    process_aug(pic_dir,xml_dir,aug_pic_dir,aug_xml_dir,aug_times)
